package com.lxf.clouduser.aop;

import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.lxf.clouduser.annotate.RateLimit;
import com.lxf.clouduser.utils.IpAddressUtils;
import com.lxf.clouduser.utils.RedisUtils;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;

/**
 * @author 罗兴发
 * @data 2025/7/1 15:35
 * @description：限制操作次数
 */
@Aspect
@Component
@Slf4j
public class RateLimitAspect {

    @Resource
    private RedisUtils redisUtils;
    @Resource
    private IpAddressUtils ipAddressUtils;

    @Resource
    private HttpServletRequest request;

    //定义黑名单key前缀
    private static final String BLACKLIST_KEY_PREFIX = "blacklist:";
    //定义白名单key前缀
    private static final String WHITELIST_KEY_PREFIX = "whitelist:";

    @Around("@annotation(rateLimit)")
    public Object rateLimit(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        JSONObject js = new JSONObject();
        //获取IP
        String ip = ipAddressUtils.getIpAddress(request);
        String key = generateKey(joinPoint, ip);
        //黑名单则直接异常
        if (isBlacklisted(ip)) {
            js.put("code",429); //设置响应状态码
            js.put("msg","Too many requests from this IP address");
            return JSONObject.toJSONString(js, SerializerFeature.WriteMapNullValue);
        }

        //如果是白名单下的不做限制
        if (isWhitelisted(ip)) {
            return joinPoint.proceed();
        }


        int limit = rateLimit.value();
        int timeout = rateLimit.timeout();

        String countStr = (String) redisUtils.get(key);
        int count = countStr == null ? 0 : Integer.parseInt(countStr);

        if (count < limit) {
            redisUtils.set(key, String.valueOf(count + 1), timeout);
            return joinPoint.proceed();
        } else {
            addToBlacklist(ip);
            js.put("code",429); //设置响应状态码
            js.put("msg","Too many requests from this IP address");
            return JSONObject.toJSONString(js, SerializerFeature.WriteMapNullValue);
        }
    }

    private boolean isBlacklisted(String ip) {
        return redisUtils.exists(BLACKLIST_KEY_PREFIX + ip);
    }

    private boolean isWhitelisted(String ip) {
        return redisUtils.exists(WHITELIST_KEY_PREFIX + ip);
    }

    private void addToBlacklist(String ip) {
        redisUtils.set(BLACKLIST_KEY_PREFIX + ip, "true", 1);
    }

    private String generateKey(ProceedingJoinPoint joinPoint, String ip) {
        String methodName = joinPoint.getSignature().getName();
        String className = joinPoint.getTarget().getClass().getName();
        return className + ":" + methodName + ":" + ip;
    }
}
